package cn.tedu;
import javax.swing.*;
import java.awt.*;
import java.awt.event.*;
import java.util.ArrayList;
import java.util.Arrays;

public class PrimGui extends JFrame {
    private final ArrayList<Point> points = new ArrayList<>();
    private final int[][] graph;
    private final int V;
    private Point startDrag, endDrag;

    public PrimGui(int vertices) {
        V = vertices;
        graph = new int[V][V];
        initUI();
    }

    private void initUI() {
        setTitle("Prim Algorithm GUI");
        setSize(800, 600);
        setDefaultCloseOperation(EXIT_ON_CLOSE);
        setLocationRelativeTo(null);

        JPanel panel = new JPanel() {
            @Override
            protected void paintComponent(Graphics g) {
                super.paintComponent(g);
                drawGraph(g);
            }
        };

        panel.addMouseListener(new MouseAdapter() {
            @Override
            public void mousePressed(MouseEvent e) {
                startDrag = findNearestPoint(e.getPoint());
                if (startDrag == null) {
                    startDrag = e.getPoint();
                    points.add(startDrag);
                    repaint();
                }
            }

            @Override
            public void mouseReleased(MouseEvent e) {
                endDrag = findNearestPoint(e.getPoint());
                if (startDrag != null && endDrag != null && startDrag != endDrag) {
                    try {
                        int weight = Integer.parseInt(JOptionPane.showInputDialog("Enter weight:"));
                        int startIndex = points.indexOf(startDrag);
                        int endIndex = points.indexOf(endDrag);
                        graph[startIndex][endIndex] = weight;
                        graph[endIndex][startIndex] = weight;
                    } catch (NumberFormatException ex) {
                        JOptionPane.showMessageDialog(null, "Invalid weight. Please enter a number.");
                    }
                    repaint();
                }
                startDrag = null;
                endDrag = null;
            }
        });

        panel.addMouseMotionListener(new MouseMotionAdapter() {
            @Override
            public void mouseDragged(MouseEvent e) {
                endDrag = e.getPoint();
                repaint();
            }
        });

        add(panel);

        JButton calculateButton = new JButton("Calculate MST");
        calculateButton.addActionListener(e -> calculateMST());
        add(calculateButton, BorderLayout.SOUTH);
    }

    private Point findNearestPoint(Point p) {
        for (Point point : points) {
            if (p.distance(point) < 10) {
                return point;
            }
        }
        return null;
    }

    private void drawGraph(Graphics g) {
        g.setColor(Color.BLACK);
        for (int i = 0; i < points.size(); i++) {
            Point p = points.get(i);
            g.fillOval(p.x - 5, p.y - 5, 10, 10);
            g.drawString(String.valueOf(i + 1), p.x - 10, p.y - 10);
        }

        g.setColor(Color.GRAY);
        for (int i = 0; i < V; i++) {
            for (int j = i + 1; j < V; j++) {
                if (graph[i][j] != 0) {
                    Point p1 = points.get(i);
                    Point p2 = points.get(j);
                    g.drawLine(p1.x, p1.y, p2.x, p2.y);
                    g.drawString(String.valueOf(graph[i][j]), (p1.x + p2.x) / 2, (p1.y + p2.y) / 2);
                }
            }
        }

        if (startDrag != null && endDrag != null) {
            g.setColor(Color.RED);
            g.drawLine(startDrag.x, startDrag.y, endDrag.x, endDrag.y);
        }
    }

    private void calculateMST() {
        int numPoints = points.size();
        if (numPoints < 2) {
            JOptionPane.showMessageDialog(this, "Not enough vertices to form a spanning tree.");
            return;
        }

        int[] parent = new int[numPoints];
        int[] key = new int[numPoints];
        boolean[] mstSet = new boolean[numPoints];
        Arrays.fill(key, Integer.MAX_VALUE);
        key[0] = 0;
        parent[0] = -1;

        for (int count = 0; count < numPoints - 1; count++) {
            int u = minKey(key, mstSet);
            mstSet[u] = true;

            for (int v = 0; v < numPoints; v++) {
                if (graph[u][v] != 0 && !mstSet[v] && graph[u][v] < key[v]) {
                    parent[v] = u;
                    key[v] = graph[u][v];
                }
            }
        }

        printMST(parent, numPoints);
    }

    private int minKey(int[] key, boolean[] mstSet) {
        int min = Integer.MAX_VALUE, minIndex = -1;

        for (int v = 0; v < key.length; v++) {
            if (!mstSet[v] && key[v] < min) {
                min = key[v];
                minIndex = v;
            }
        }
        return minIndex;
    }

    private void printMST(int[] parent, int numPoints) {
        StringBuilder sb = new StringBuilder();
        sb.append("Edge\tWeight\n");
        for (int i = 1; i < numPoints; i++) {
            if (parent[i] != -1) {
                sb.append((parent[i] + 1)).append(" - ").append((i + 1)).append("\t").append(graph[i][parent[i]]).append("\n");
            }
        }
        JOptionPane.showMessageDialog(this, sb.toString(), "Minimum Spanning Tree", JOptionPane.INFORMATION_MESSAGE);
    }

    public static void main(String[] args) {
        SwingUtilities.invokeLater(() -> {
            PrimGui ex = new PrimGui(8);  // Adjust the number of vertices here
            ex.setVisible(true);
        });
    }
}
